#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
投资策略生成模块
根据股票趋势分析结果生成每日投资策略
"""

import os
import json
import yaml
import logging
import datetime
import pandas as pd
import numpy as np
from pathlib import Path
from collections import defaultdict

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class StrategyGenerator:
    """投资策略生成类，负责根据股票趋势分析结果生成投资策略"""
    
    def __init__(self, config_path='config/config.yaml'):
        """
        初始化策略生成器
        
        Args:
            config_path: 配置文件路径
        """
        self.config = self._load_config(config_path)
        self.prediction_config = self.config['prediction']
        self.stock_market_config = self.config['stock_market']
        self.results_path = self.config['system']['data_storage']['results_path']
        
        # 确保存储目录存在
        Path(self.results_path).mkdir(parents=True, exist_ok=True)
        
        # 置信度阈值
        self.confidence_threshold = self.prediction_config['settings']['confidence_threshold']
        
        logger.info("策略生成器初始化完成")
    
    def _load_config(self, config_path):
        """
        加载配置文件
        
        Args:
            config_path: 配置文件路径
            
        Returns:
            dict: 配置信息
        """
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
            logger.info(f"成功加载配置文件: {config_path}")
            return config
        except Exception as e:
            logger.error(f"加载配置文件失败: {e}")
            raise
    
    def load_analysis_results(self, analysis_file):
        """
        加载股票分析结果
        
        Args:
            analysis_file: 股票分析结果文件路径
            
        Returns:
            list: 股票分析结果列表
        """
        try:
            with open(analysis_file, 'r', encoding='utf-8') as f:
                analysis_results = json.load(f)
            logger.info(f"成功加载 {len(analysis_results)} 条股票分析结果")
            return analysis_results
        except Exception as e:
            logger.error(f"加载股票分析结果失败: {e}")
            return []
    
    def filter_results(self, analysis_results):
        """过滤分析结果，只保留有效结果"""
        if isinstance(analysis_results, list):
            # 处理列表结果
            filtered_results = {}
            for result in analysis_results:
                stock_code = result.get('stock_code')
                prediction_data = result.get('prediction', {})
                
                # 安全地处理预测数据
                if isinstance(prediction_data, dict):
                    confidence = prediction_data.get('confidence', 0)
                    trend = prediction_data.get('trend', 'unknown')
                    
                    if confidence >= self.confidence_threshold:
                        filtered_results[stock_code] = result
                else:
                    # 预测数据不是字典的情况，忽略这个结果
                    logger.warning(f"Stock {stock_code} prediction is not a dictionary: {prediction_data}")
            
            return filtered_results
        else:
            # 处理字典结果
            filtered_results = {}
            for stock_code, result in analysis_results.items():
                prediction_data = result.get('prediction', {})
                
                # 安全地处理预测数据
                if isinstance(prediction_data, dict):
                    confidence = prediction_data.get('confidence', 0)
                    trend = prediction_data.get('trend', 'unknown')
                    
                    if confidence >= self.confidence_threshold:
                        filtered_results[stock_code] = result
                else:
                    # 预测数据不是字典的情况，忽略这个结果
                    logger.warning(f"Stock {stock_code} prediction is not a dictionary: {prediction_data}")
            
            return filtered_results
    
    def categorize_stocks(self, filtered_results):
        """
        根据预测趋势对股票进行分类
        
        Args:
            filtered_results: 过滤后的分析结果字典
            
        Returns:
            dict: 分类结果，格式为 {'up': [], 'down': [], 'stable': []}
        """
        categorized = {
            'up': [],
            'down': [],
            'stable': []
        }
        
        # Handle filtered_results being a dictionary (stock_code -> result)
        if isinstance(filtered_results, dict):
            for stock_code, result in filtered_results.items():
                # 确保结果包含prediction字段且为字典
                prediction = result.get('prediction', {})
                if not isinstance(prediction, dict):
                    logger.warning(f"Stock {stock_code} prediction is not a dictionary: {prediction}")
                    continue
                    
                # 获取预测趋势
                trend = prediction.get('trend', 'unknown')
                
                # 分类
                if trend in categorized:
                    # 添加stock_code到结果中，确保后续处理可以访问
                    if 'stock_code' not in result:
                        result['stock_code'] = stock_code
                    categorized[trend].append(result)
        # Handle filtered_results being a list of results
        elif isinstance(filtered_results, list):
            for result in filtered_results:
                # 获取预测趋势
                prediction = result.get('prediction', {})
                if not isinstance(prediction, dict):
                    stock_code = result.get('stock_code', 'unknown')
                    logger.warning(f"Stock {stock_code} prediction is not a dictionary: {prediction}")
                    continue
                    
                trend = prediction.get('trend', 'unknown')
                
                # 分类
                if trend in categorized:
                    categorized[trend].append(result)
        
        # 按照置信度排序
        for trend in categorized:
            categorized[trend] = sorted(
                categorized[trend],
                key=lambda x: x.get('prediction', {}).get('confidence', 0) 
                    if isinstance(x.get('prediction', {}), dict) else 0,
                reverse=True
            )
        
        logger.info(f"股票分类结果: 上涨={len(categorized['up'])}, 下跌={len(categorized['down'])}, 稳定={len(categorized['stable'])}")
        return categorized
    
    def analyze_industry_trends(self, filtered_results):
        """
        分析行业趋势
        
        Args:
            filtered_results: 过滤后的分析结果(字典或列表)
            
        Returns:
            dict: 行业趋势分析结果
        """
        # 行业趋势统计
        industry_trends = defaultdict(lambda: {'up': 0, 'down': 0, 'stable': 0, 'stocks': []})
        
        # 处理字典格式的结果 {stock_code: result}
        if isinstance(filtered_results, dict):
            for stock_code, result in filtered_results.items():
                # 确保预测是字典类型
                prediction = result.get('prediction', {})
                if not isinstance(prediction, dict):
                    logger.warning(f"Stock {stock_code} prediction is not a dictionary: {prediction}")
                    continue
                
                # 获取预测趋势
                trend = prediction.get('trend', 'unknown')
                
                # 获取行业信息（这里可能需要从result中获取，或通过其他方式）
                industry = result.get('industry', 'unknown')
                
                # 使用默认行业如果未找到
                if industry == 'unknown':
                    # 简化处理，假设行业信息在result中或通过其他方式获取
                    for ind in self.stock_market_config['industries']:
                        # 简化处理，随机分配行业
                        if hash(stock_code) % len(self.stock_market_config['industries']) == self.stock_market_config['industries'].index(ind):
                            industry = ind
                            break
                
                # 更新行业趋势
                if trend in ['up', 'down', 'stable']:
                    industry_trends[industry][trend] += 1
                
                # 添加股票到行业
                industry_trends[industry]['stocks'].append({
                    'stock_code': stock_code,
                    'trend': trend,
                    'confidence': prediction.get('confidence', 0)
                })
        # 处理列表格式的结果 [result1, result2, ...]
        elif isinstance(filtered_results, list):
            for result in filtered_results:
                # 获取股票代码和预测趋势
                stock_code = result.get('stock_code', '')
                
                # 确保预测是字典类型
                prediction = result.get('prediction', {})
                if not isinstance(prediction, dict):
                    logger.warning(f"Stock {stock_code} prediction is not a dictionary: {prediction}")
                    continue
                
                # 获取预测趋势
                trend = prediction.get('trend', 'unknown')
                
                # 获取行业信息（这里可能需要从result中获取，或通过其他方式）
                industry = result.get('industry', 'unknown')
                
                # 使用默认行业如果未找到
                if industry == 'unknown':
                    # 简化处理，通过股票代码分配行业
                    for ind in self.stock_market_config['industries']:
                        # 简化处理，随机分配行业
                        if hash(stock_code) % len(self.stock_market_config['industries']) == self.stock_market_config['industries'].index(ind):
                            industry = ind
                            break
                
                # 更新行业趋势
                if trend in ['up', 'down', 'stable']:
                    industry_trends[industry][trend] += 1
                
                # 添加股票到行业
                industry_trends[industry]['stocks'].append({
                    'stock_code': stock_code,
                    'trend': trend,
                    'confidence': prediction.get('confidence', 0)
                })
        
        # 计算每个行业的整体趋势
        industry_summary = {}
        
        for industry, data in industry_trends.items():
            total = data['up'] + data['down'] + data['stable']
            if total > 0:
                up_ratio = data['up'] / total
                down_ratio = data['down'] / total
                
                if up_ratio > 0.6:
                    overall_trend = 'up'
                elif down_ratio > 0.6:
                    overall_trend = 'down'
                else:
                    overall_trend = 'mixed'
                
                industry_summary[industry] = {
                    'trend': overall_trend,
                    'up_ratio': up_ratio,
                    'down_ratio': down_ratio,
                    'stock_count': total,
                    'stocks': data['stocks']
                }
        
        logger.info(f"完成 {len(industry_summary)} 个行业的趋势分析")
        return industry_summary
    
    def generate_investment_advice(self, categorized_stocks, industry_trends):
        """
        基于分类股票和行业趋势生成投资建议
        
        Args:
            categorized_stocks: 分类的股票字典
            industry_trends: 行业趋势分析结果
            
        Returns:
            dict: 投资建议
        """
        logger.info("开始生成投资建议")
        
        try:
            investment_advice = {
                'buy': [],
                'sell': [],
                'hold': [],
                'watch': [],
                'industry_recommendations': [],
                'market_summary': {
                    'trend': 'neutral',
                    'confidence': 0,
                    'recommendation': '市场整体表现平稳，建议适度投资。'
                }
            }
            
            # 处理分类的股票
            if not categorized_stocks:
                logger.warning("没有分类的股票信息")
                return investment_advice
            
            # 确保分类股票是预期的字典结构
            if not isinstance(categorized_stocks, dict):
                logger.warning(f"分类股票不是预期的字典结构: {type(categorized_stocks)}")
                return investment_advice
                
            # 生成股票推荐
            for trend, stocks in categorized_stocks.items():
                if not stocks:
                    continue
                    
                for stock_info in stocks:
                    if not isinstance(stock_info, dict):
                        logger.warning(f"股票信息不是预期的字典结构: {type(stock_info)}")
                        continue
                        
                    stock_code = stock_info.get('stock_code', '')
                    confidence = stock_info.get('confidence', 0)
                    
                    if not stock_code:
                        continue
                    
                    # 根据趋势和置信度生成推荐
                    recommendation = {
                        'stock_code': stock_code,
                        'confidence': confidence,
                        'reason': f"基于{confidence:.2f}的置信度预测趋势为{trend}"
                    }
                    
                    # 添加股票到相应的推荐列表
                    if trend == 'up' and confidence >= self.confidence_threshold:
                        investment_advice['buy'].append(recommendation)
                    elif trend == 'down' and confidence >= self.confidence_threshold:
                        investment_advice['sell'].append(recommendation)
                    elif trend == 'stable' and confidence >= self.confidence_threshold:
                        investment_advice['hold'].append(recommendation)
                    else:
                        investment_advice['watch'].append(recommendation)
            
            # 处理行业趋势
            if industry_trends:
                for industry, trend_info in industry_trends.items():
                    if not isinstance(trend_info, dict):
                        continue
                        
                    industry_trend = trend_info.get('trend', 'neutral')
                    industry_confidence = trend_info.get('confidence', 0)
                    industry_stocks = trend_info.get('stocks', [])
                    
                    # 生成行业推荐
                    industry_recommendation = {
                        'industry': industry,
                        'trend': industry_trend,
                        'confidence': industry_confidence,
                        'stocks': industry_stocks[:5],  # 最多展示5只股票
                        'recommendation': self._generate_industry_recommendation(industry, industry_trend, industry_confidence)
                    }
                    
                    investment_advice['industry_recommendations'].append(industry_recommendation)
            
            # 生成市场总结
            market_summary = self._generate_market_summary(categorized_stocks, industry_trends)
            if market_summary:
                investment_advice['market_summary'] = market_summary
            
            logger.info("成功生成投资建议")
            return investment_advice
            
        except Exception as e:
            import traceback
            logger.error(f"生成投资建议失败: {e}")
            logger.error(traceback.format_exc())
            return {
                'buy': [],
                'sell': [],
                'hold': [],
                'watch': [],
                'industry_recommendations': [],
                'market_summary': {
                    'trend': 'neutral',
                    'confidence': 0,
                    'recommendation': '数据处理错误，无法给出准确建议。'
                }
            }
    
    def _generate_industry_recommendation(self, industry, trend, confidence):
        """
        生成行业推荐
        
        Args:
            industry: 行业名称
            trend: 行业趋势
            confidence: 行业置信度
            
        Returns:
            str: 行业推荐
        """
        # 实现行业推荐的逻辑
        # 这里可以根据行业和趋势生成推荐语句
        return f"基于{confidence:.2f}的置信度预测行业{industry}的趋势为{trend}"
    
    def _generate_market_summary(self, categorized_stocks, industry_trends):
        """
        生成市场总结
        
        Args:
            categorized_stocks: 分类的股票字典
            industry_trends: 行业趋势分析结果
            
        Returns:
            dict: 市场总结
        """
        # 实现市场总结的逻辑
        # 这里可以根据股票分类和行业趋势生成市场总结
        return {
            'trend': 'neutral',
            'confidence': 0,
            'recommendation': '市场整体表现平稳，建议适度投资。'
        }
    
    def generate_strategy_report(self, investment_advice):
        """
        生成策略报告
        
        Args:
            investment_advice: 投资建议
            
        Returns:
            str: 策略报告文本
        """
        # 获取当前日期
        today = datetime.datetime.now().strftime('%Y-%m-%d')
        
        # 市场趋势
        market_trend = investment_advice['market_summary']['trend']
        market_trend_desc = {
            'bullish': '强烈看涨',
            'slightly_bullish': '略微看涨',
            'neutral': '中性',
            'slightly_bearish': '略微看跌',
            'bearish': '强烈看跌',
            'unknown': '未知'
        }.get(market_trend, '未知')
        
        # 行业推荐
        industry_recommendations = investment_advice['industry_recommendations']
        industry_text = ""
        for recommendation in industry_recommendations:
            industry_text += f"- {recommendation['industry']}：趋势 {recommendation['trend']}，置信度 {recommendation['confidence']:.2f}\n"
        
        # 市场总结
        market_summary = investment_advice['market_summary']
        market_summary_text = market_summary['recommendation']
        
        # 生成报告
        report = f"""
# 每日投资策略报告 ({today})

## 市场概览

- **市场趋势**：{market_trend_desc}

## 行业分析

{industry_text if industry_text else "暂无行业分析\n"}

## 市场总结

{market_summary_text}

## 风险提示

本报告基于新闻分析和技术指标生成，仅供参考，不构成投资建议。投资有风险，入市需谨慎。
        """
        
        logger.info("成功生成策略报告")
        return report
    
    def save_strategy_report(self, report):
        """
        保存策略报告
        
        Args:
            report: 策略报告文本
            
        Returns:
            str: 保存的文件路径
        """
        try:
            # 生成文件名
            timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f"investment_strategy_{timestamp}.md"
            filepath = os.path.join(self.results_path, filename)
            
            # 保存为Markdown文件
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(report)
            
            logger.info(f"成功保存策略报告到 {filepath}")
            return filepath
        
        except Exception as e:
            logger.error(f"保存策略报告失败: {e}")
            return None
    
    def save_investment_advice(self, investment_advice):
        """
        保存投资建议
        
        Args:
            investment_advice: 投资建议
            
        Returns:
            str: 保存的文件路径
        """
        try:
            # 生成文件名
            timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f"investment_advice_{timestamp}.json"
            filepath = os.path.join(self.results_path, filename)
            
            # 保存为JSON文件
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(investment_advice, f, ensure_ascii=False, indent=2)
            
            logger.info(f"成功保存投资建议到 {filepath}")
            return filepath
        
        except Exception as e:
            logger.error(f"保存投资建议失败: {e}")
            return None
    
    def generate_strategy(self, analysis_results=None, report_file=None):
        """
        根据分析结果生成投资策略
        
        Args:
            analysis_results: 分析结果字典或分析结果文件路径
            report_file: 策略报告保存路径
            
        Returns:
            investment_advice: 投资建议
            strategy_report: 策略报告
            report_path: 策略报告保存路径
        """
        try:
            # 判断输入类型和有效性
            if not analysis_results:
                logger.warning("没有找到分析结果或结果为空")
                return None, None, None
                
            if isinstance(analysis_results, str):
                # 输入是文件路径，加载分析结果
                logger.info(f"从文件加载分析结果: {analysis_results}")
                analysis_data = self.load_analysis_results(analysis_results)
                if not analysis_data:
                    logger.warning("加载的分析结果为空")
                    return None, None, None
            else:
                # 直接使用结果字典或列表
                logger.info(f"使用直接传入的分析结果，包含 {len(analysis_results)} 条记录")
                analysis_data = analysis_results
            
            # Debug logging to help identify the prediction structure
            if isinstance(analysis_data, dict) and analysis_data:
                for stock_code, result in list(analysis_data.items())[:1]:  # Just check first item
                    logger.info(f"Debug - stock_code: {stock_code}")
                    logger.info(f"Debug - result keys: {list(result.keys())}")
                    prediction = result.get('prediction', {})
                    logger.info(f"Debug - prediction type: {type(prediction)}")
                    
                    if isinstance(prediction, dict):
                        logger.info(f"Debug - prediction keys: {list(prediction.keys())}")
                    else:
                        logger.info(f"Debug - prediction is not a dict: {prediction}")
            
            # 过滤分析结果，只保留有效结果
            filtered_results = self.filter_results(analysis_data)
            
            if not filtered_results:
                logger.warning("过滤后没有有效的分析结果")
                return None, None, None
            
            # 分类股票
            categorized_stocks = self.categorize_stocks(filtered_results)
            
            # 分析行业趋势
            industry_trends = self.analyze_industry_trends(filtered_results)
            
            # 生成投资建议
            investment_advice = self.generate_investment_advice(categorized_stocks, industry_trends)
            
            # 生成策略报告
            strategy_report = self.generate_strategy_report(investment_advice)
            
            # 保存投资建议和策略报告
            advice_path = self.save_investment_advice(investment_advice)
            report_path = self.save_strategy_report(strategy_report)
            
            logger.info("成功生成投资策略")
            return investment_advice, strategy_report, report_path
            
        except Exception as e:
            import traceback
            logger.error(f"生成投资策略失败: {e}")
            logger.error(traceback.format_exc())
            return None, None, None

def main():
    """主函数"""
    import sys
    
    try:
        # 检查命令行参数
        if len(sys.argv) < 2:
            print("用法: python strategy_generator.py <analysis_results_file_path>")
            return
        
        analysis_file = sys.argv[1]
        
        # 创建策略生成器
        generator = StrategyGenerator()
        
        # 生成投资策略
        investment_advice, strategy_report, report_path = generator.generate_strategy(analysis_file)
        
        if report_path:
            print(f"成功生成投资策略并保存到 {report_path}")
        else:
            print("生成投资策略失败")
        
    except Exception as e:
        logger.error(f"策略生成失败: {e}")

if __name__ == "__main__":
    main() 